package com.javatest.springboot2ratelimit.aop;

import com.google.common.collect.ImmutableList;
import com.javatest.springboot2ratelimit.annoation.MakeRateLimit;
import com.javatest.springboot2ratelimit.enums.RateLimitFineness;
import com.javatest.springboot2ratelimit.exception.RateLimitException;
import java.io.Serializable;
import java.lang.reflect.Method;
import java.util.Random;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author lzy
 * @create 2020-07-21 下午3:14
 */
@Aspect
//@Component
public class RateLimitByRedisCountAspect extends AbstractRateLimitAspect {
    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitByRedisCountAspect.class);
    private final String REDIS_SCRIPT = buildLuaScript();
    @Autowired
    private RedisTemplate<String, Serializable> redisTemplate;

    @Before(value = "limitPointCut() && @annotation(rateLimit)", argNames = "joinPoint, rateLimit")
    @Override
    public void before(JoinPoint joinPoint, MakeRateLimit rateLimit) {
        final Signature signature = joinPoint.getSignature();
        // 完整类名
        final String className = signature.getDeclaringType().getName();
        // 注解所在方法名
        Method method = ((MethodSignature) joinPoint.getSignature()).getMethod();
        String methodName = method.getName();

        // 根据不同限流策略控制限流器的配置粒度
        // 默认采用类名_方法名作为key
        final String prefix = className + "_" + methodName;
        String key = prefix;
        if (rateLimit.limitFineness() == RateLimitFineness.IP) {
            key += "_" + getIpAddress();
        } else if (rateLimit.limitFineness() == RateLimitFineness.User) {
            // todo: 这里用一个随机数模拟用户id
            key += "_" + new Random(10).nextInt(5);
        }

        ImmutableList<String> keys = ImmutableList.of(key);
        RedisScript<Number> redisScript = new DefaultRedisScript<>(REDIS_SCRIPT, Number.class);
        Number count = redisTemplate.execute(redisScript, keys, rateLimit.limitCount(), rateLimit.limitPeriod());
        if (!(count != null && count.intValue() <= rateLimit.limitCount())) {
            LOGGER.info("[Redis-count] limiter-{} 对接口访问进行限速", key);
            throw new RateLimitException("访问过于频繁, 拒绝了请求!");
        }
    }

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    private String buildLuaScript() {
        String lua = "local c" +
                "\nc = redis.call('get', KEYS[1])" +
                // 调用不超过最大值，则直接返回
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                // 执行计算器自加
                "\nc = redis.call('incr', KEYS[1])" +
                "\nif tonumber(c) == 1 then" +
                // 从第一次调用开始限流，设置对应键值的过期
                "\nredis.call('expire', KEYS[1], ARGV[2])" +
                "\nend" +
                "\nreturn c;";
        return lua;
    }

}
